# -*- coding:utf-8 -*-
""""
@file name  : weight_fmap_visualization.py
@author     : QuZhang
@date       : 2020-12-24 22:30
@brief      : 可视化卷积核和特征图
"""
from tools.common_tools import set_seed
from torch.utils.tensorboard import SummaryWriter
import torchvision.models as models
import torch
import torch.nn as nn
import torchvision.utils as vutils
from torchvision.transforms import transforms
from PIL import Image


set_seed(1)

# ----------- kernel visualization ----------
flag = True
if flag:
    writer = SummaryWriter(comment='kernel_visualizaton', filename_suffix='_test_kernel')
    alexnet = models.alexnet(pretrained=True)  # 加载预训练模型

    kernel_num = -1  # 指示当前是第几个卷积层
    vis_max = 1  # 最大可视化层,也就是可视化0和1层

    with torch.no_grad():
        # 遍历所有的层
        for sub_module in alexnet.modules():
            # 卷积层
            if isinstance(sub_module, nn.Conv2d):
                kernel_num += 1
                if kernel_num > vis_max:
                    break
                kernels = sub_module.weight  # 读取改层的所有卷积核
                c_out, c_int, k_h, k_w = tuple(kernels.shape)  # (卷积核个数，一个卷积核的通道数，一个通道的张量尺寸)

                # 1.可视化每一个核里的每一个通道图
                # 遍历该卷积核的每一个核
                # 一个核显示在一行，一个核的一个通道为一列
                for o_idx in range(c_out):
                    kernel_idx = kernels[o_idx, :, :, :].unsqueeze(1)  # 获取第o_idx+1个核
                    kernel_grid = vutils.make_grid(kernel_idx, normalize=True, scale_each=True, nrow=c_int)  # 每一个核的通道数为c_int
                    writer.add_image('{}_Convlayer_split_in_channel'.format(kernel_num), kernel_grid, global_step=o_idx)

                # 2.将每一个核变为3通道后可视化
                # tensor.view()改变张量的形状，view的参数就是改变后张量的形状
                kernel_all = kernels.view(-1, 3, k_h, k_w)  # 把改层的每一个核都变成(3，h，w)
                kernel_grid = vutils.make_grid(kernel_all, normalize=True, scale_each=True, nrow=8)
                writer.add_image("{}_all".format(kernel_num), kernel_grid, global_step=322)

                print("{}_convlayer shape:{}".format(kernel_num, tuple(kernels.shape)))

        writer.close()

flag = True
if flag:
    with torch.no_grad():
        writer = SummaryWriter(comment="lena_visualization", filename_suffix="lena_kernel")

        # 数据
        path_img = "./lena.png"
        normMean = [0.49139968, 0.48215827, 0.44653124]
        normStd = [0.24703233, 0.24348505, 0.26158768]

        # 数据增强
        img_transforms = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=normMean, std=normStd)
        ])

        img_pil = Image.open(path_img).convert("RGB")  # 1.读取图片
        if img_transforms is not None:
            img_tensor = img_transforms(img_pil)  # 2. 图像增强
        img_tensor.unsqueeze_(0)  # 模型需要四维张量, chw -> bchw

        # 模型
        alexnet = models.alexnet(pretrained=True)

        # forward
        convlayer1 = alexnet.features[0]
        fmap_1 = convlayer1(img_tensor)

        # 预处理,显示输出特征图的每一个通道
        fmap_1.transpose_(0, 1)  # 交换维度: bchw=(1, 64, 55, 55) -> (64, 1, 55, 55)

        # print('before shape: ', fmap_1.shape)
        # fmap_1.squeeze_()
        # print('after shape: ', fmap_1.shape)

        fmap_1_grid = vutils.make_grid(fmap_1, normalize=True, scale_each=True, nrow=8)
        writer.add_image('feature map in conv1', fmap_1_grid, global_step=322)
        writer.close()
